import os

from PIL import Image
from matplotlib import pyplot as plt
from torchvision import transforms
from tqdm import tqdm


def read_PIL(image_path):
    # 使用PIL打开图像，方便之后使用该库中的图像增强
    image = Image.open(image_path)
    return image.convert("RGB")


def getdirs(dir):
    filenames = os.listdir(dir)
    return filenames


def getfiles(dir):
    filenames = os.listdir(dir)
    return filenames


def img_intensive(dir, batch_size, onlyTrain):
    data_train = dir + "train\\"
    data_test = dir + "test\\"
    if onlyTrain:
        generate_img(batch_size, data_train)
    else:
        generate_img(batch_size, data_train)
        generate_img(batch_size, data_test)
    print("图像增强成功")


def generate_img(batch_size, data_train):
    for i in getdirs(data_train):
        print("当前执行文件夹名称:{}".format(i))
        img_dir = data_train + i
        filename = getfiles(img_dir)
        for j in tqdm(filename):
            img = read_PIL(img_dir + "\\" + j)
            for m in range(1, batch_size):
                data = transforms.Compose([transforms.RandomRotation(45),
                                           transforms.CenterCrop(224),
                                           transforms.RandomHorizontalFlip(p=0.5),
                                           transforms.RandomVerticalFlip(p=0.5),
                                           transforms.ColorJitter(brightness=0.2, contrast=0.1, saturation=0.1,
                                                                  hue=0.1),
                                           transforms.RandomGrayscale(p=0.025),
                                           transforms.ToTensor()
                                           ])(img)
                plt.imshow(data.permute(1, 2, 0))
                plt.axis('off')
                plt.savefig('{}/intensive_{}_{}.png'.format(img_dir, j.split(".")[0], m),
                            bbox_inches="tight",
                            pad_inches=0.0)
